#!/bin/bash
#set -ex
set -e
# -------------------------------------------------------------------------------
# Filename:     benchmark4vllm_72b.sh
# UpdateDate:   2024/08/15
# Description:  一键运行性能测试，并保存性能测试日志及测试期间的cnmon信息到各自log。
# Example:      ./benchmark4vllm_72b.sh
# Depends:
# Notes:        本脚本适用于docker容器环境下运行；
# -------------------------------------------------------------------------------
CMD_TIME=$(date +%Y%m%d%H%M%S.%N)
# 进入工作目录
CUR_DIR=$(cd $(dirname $0);pwd)
pushd "${CUR_DIR}/.."
if [ ! -d "./log" ]; then mkdir -p "./log";fi
popd
#每次benchmark的sleep时间,长一些会更好
TIME_SLEEP=60
#export MLU_VISIBLE_DEVICES=4
export VLLM_LATENCY_DEBUG=true
export CN_TASKTOPO_RESIDENT=FALSE
#export USE_PAGED=false
#设置模型路径
#MODEL_Name=qwen2-72b
MODEL_FULLNAME=${MODEL_PATH}
#export MAX_TOTAL_TOKENS=512
block_size=16
max_model_len=4096
max_num_seqs=256
array_bs=(1 2 4)
seqlen_output_bs=(256)

for seqlen_input in 128 256 512 1024 2048; do
#for seqlen_input in 1024 2048; do
    #不同条件不同赋值测试,减少部分测试
    case $seqlen_input in
        128)
            seqlen_output_bs=(128 256)
            array_bs=(1500)     #gpu_memory_utilization=0.98, 1536:1444
            ;;
        256)
            seqlen_output_bs=(256 512)
            array_bs=(240 256)      #gpu_memory_utilization=0.98, 1024:244
            ;;
        512)
            seqlen_output_bs=(256 512)
            array_bs=(120 128)      #gpu_memory_utilization=0.98, 512:122
            ;;
        1024)
            seqlen_output_bs=(256 1024)
            array_bs=(64)       #gpu_memory_utilization=0.98, 256:61
            ;;
        2048)
            seqlen_output_bs=(256 2048)
            array_bs=(32)       #gpu_memory_utilization=0.98, 128:30
            ;;
        4096)
            seqlen_output_bs=(4096)
            array_bs=(16)       #gpu_memory_utilization=0.98,
            ;;
        *)
            seqlen_output_bs=(256)
            array_bs=(1 2 4)
            ;;
    esac
    for seqlen_output in "${seqlen_output_bs[@]}"; do
        for bs in "${array_bs[@]}"; do
            for tp_nums in ${MLU_TP_NUM}; do
                # 0. 设置参数
                #export MAX_TOTAL_TOKENS=$(expr ${seqlen_input} + ${seqlen_output})
                #block_size=${MAX_TOTAL_TOKENS}
                # 1>. set max_num_seqs。在1个推理阶段中，LLMEngine最多能处理的seq数量（1条seq就是指我们待推理的1条数据）。默认是256
                if [ $bs -gt 256 ]; then max_num_seqs=$bs; else max_num_seqs=256; fi
                # 2>. set max_model_len . 每个模型不同，根据模型的config.json里的参数定，比如llama2-7b是4096。# max_model_len=4096
                max_model_len=$(expr ${seqlen_input} + ${seqlen_output})
                # 3>. max_num_batched_tokens. 在1个推理阶段中，LLMEngine最多能处理的token数量。默认是2048
                max_num_batched_tokens=$(expr ${seqlen_input} \* $bs)
                if [ $max_model_len -gt $max_num_batched_tokens ]; then max_num_batched_tokens=$max_model_len; fi
                # max_num_batched_tokens must be greater than or equal to max_num_seqs.
                if [ $max_num_seqs -gt $max_num_batched_tokens ]; then max_num_batched_tokens=$max_num_seqs; fi

                # 1. 监控 cnmon 命令并将输出追加到日志文件中
                LOG_FILENAME="${CUR_DIR}/../log/${MODEL_Name}_latency_blks${block_size}_sli${seqlen_input}_slo${seqlen_output}_tp${tp_nums}_bs${bs}_${CMD_TIME}"
                while true; do cnmon -c ${MLU_VISIBLE_DEVICES} >> "${LOG_FILENAME}_cnmon.log";sleep 1;done 2>&1 &
                PID_CNMONProcess=$!
                echo "######################################"
                #如果脚本强制或意外中断，可通过【ps -a | grep benchmark4vllm】命令或通过如下变量【PID_CNMONProcess】，查找后台运行的进行使用【sudo kill -9 $PID】结束后台进程。
                echo "PID_CNMONProcess: $PID_CNMONProcess"
                echo "LOG_FILENAME: $LOG_FILENAME"
                echo "######################################"

                # 2. 性能测试
                cd /workspace/vllm_mlu
                python benchmarks/benchmark_latency.py \
                        --model ${MODEL_FULLNAME} \
                        --tokenizer ${MODEL_FULLNAME} \
                        --num-iters-warmup 1 \
                        --num-iters 1 \
                        --dtype float16 \
                        --input-len ${seqlen_input} \
                        --output-len ${seqlen_output} \
                        --block-size ${block_size} \
                        -tp ${tp_nums} \
                        --trust-remote-code \
                        --max-model-len ${max_model_len} \
                        --max-num-batched-tokens ${max_num_batched_tokens} \
                        --max-num-seqs ${max_num_seqs} \
                        --batch-size ${bs} 2>&1 | tee "${LOG_FILENAME}.log"
##--enforce-eager \
                # 3. 删除【实时记录cnmon信息】的进程;# 使用 tail -f ./*_cnmon.log 可以实时查看日志文件的内容
                sleep 2
                kill -9 $PID_CNMONProcess
                # 4. sleep
                sleep ${TIME_SLEEP}
            done
        done
    done
done
sleep 1
#切换目录
pushd "${CUR_DIR}/.." && ls -la log/